import numpy as np
from astropy.io import fits
from astropy.time import Time
import os
import matplotlib.pyplot as plt
from datetime import timedelta
from matplotlib import gridspec
# ================================
# Step 1: FITS文件读取与数据提取
# ================================
import numpy as np
from astropy.io import fits

def load_fits_file(file_name, sample_begin, sample_end, reverse_flag=False, bit_1=False):
    """
    加载 .fits 数据，支持频率翻转和 1-bit 降位。

    参数：
        file_name     - FITS 文件路径
        sample_begin  - 起始 sample
        sample_end    - 结束 sample
        reverse_flag  - 是否频率轴翻转（高→低）
        bit_1         - 是否将强度转换为 1-bit（二值化）

    返回：
        data.T: shape = [nchan, nsample]
    """
    try:
        import fitsio
        _data, h = fitsio.read(file_name, header=True)
    except ImportError:
        with fits.open(file_name) as file:
            h = file[1].header
            _data = file[1].data

    _data = _data['DATA'].reshape(
        h['NAXIS2'] * h['NSBLK'], h['NPOL'], h['NCHAN']
    )[sample_begin:sample_end, :2, :]

    if reverse_flag:
        _data = _data[:, :, ::-1]

    # 两极化合成为强度
    scale = np.mean(_data[:, 0, :]) / np.mean(_data[:, 1, :])
    data = (_data[:, 0, :] + _data[:, 1, :] * scale) / 2

    # 1-bit 降位处理（可选）
    if bit_1:
        threshold = np.mean(data)
        data = (data > threshold).astype(np.uint8)

    return data.T


def get_total_samples(fits_file):
    with fits.open(fits_file) as hdul:
        h = hdul[1].header
        return h['NSBLK'] * h['NAXIS2']


def get_fits_freq_time_params(fits_file):
    with fits.open(fits_file) as hdul:
        hdr_primary = hdul[0].header
        hdr = hdul[1].header
        data = hdul[1].data

        tsamp = hdr['TBIN']
        freq_array = data['DAT_FREQ'][0][::-1]  # 反转频率轴
        f_low = np.min(freq_array)
        f_high = np.max(freq_array)

        # ==== MJD 推算 ====
        if 'STT_IMJD' in hdr_primary and 'STT_SMJD' in hdr_primary:
            imjd = int(hdr_primary['STT_IMJD'])
            smjd = int(hdr_primary['STT_SMJD'])
            offs = float(hdr_primary.get('STT_OFFS', 0.0))  # 有可能没有 STT_OFFS
            mjd_start = imjd + (smjd + offs) / 86400.0
        else:
            raise ValueError("！！！无法找到 STT_IMJD / STT_SMJD 来推算起始 MJD")

    return f_low, f_high, tsamp, freq_array, mjd_start


# ================================
# Step 2: 多文件拼接数据段提取
# ================================
def calculate_dm_delay_samples(DM, f_low, f_high, tsamp):
    delay_s = 4.15e3 * DM * (1.0 / f_low**2 - 1.0 / f_high**2)
    return int(np.ceil(delay_s / tsamp))


# def read_padded_samples(fits_list, current_index, start_sample, end_sample):
#     all_data = []
#     total_samples = get_total_samples(fits_list[current_index])

#     # 前补
#     if start_sample < 0:
#         prev_index = current_index - 1
#         if prev_index >= 0:
#             prev_total = get_total_samples(fits_list[prev_index])
#             need_from_prev = -start_sample
#             prev_start = max(0, prev_total - need_from_prev)
#             data_prev = load_fits_file(fits_list[prev_index], prev_start, prev_total, reverse_flag=True)
#             all_data.append(data_prev)
#         start_sample = 0

#     # 当前
#     this_end = min(end_sample, total_samples)
#     if start_sample < this_end:
#         data_this = load_fits_file(fits_list[current_index], start_sample, this_end, reverse_flag=True)
#         all_data.append(data_this)

#     # 后补
#     if end_sample > total_samples:
#         next_index = current_index + 1
#         if next_index < len(fits_list):
#             extra = end_sample - total_samples
#             data_next = load_fits_file(fits_list[next_index], 0, extra, reverse_flag=True)
#             all_data.append(data_next)

#     return np.hstack(all_data)


def read_padded_samples(fits_list, current_index, start_sample, end_sample):
    all_data = []
    total_samples = get_total_samples(fits_list[current_index])

    # 前补
    if start_sample < 0:
        prev_index = current_index - 1
        need_from_prev = -start_sample
        if prev_index >= 0:
            prev_total = get_total_samples(fits_list[prev_index])
            prev_start = max(0, prev_total - need_from_prev)
            data_prev = load_fits_file(fits_list[prev_index], prev_start, prev_total, reverse_flag=True)
        else:
            # 补零
            nchan = get_nchan_from_fits(fits_list[current_index])
            data_prev = np.zeros((nchan, need_from_prev), dtype=np.float32)
        all_data.append(data_prev)
        start_sample = 0

    # 当前文件数据
    this_end = min(end_sample, total_samples)
    if start_sample < this_end:
        data_this = load_fits_file(fits_list[current_index], start_sample, this_end, reverse_flag=True)
        all_data.append(data_this)

    # 后补
    if end_sample > total_samples:
        extra = end_sample - total_samples
        next_index = current_index + 1
        if next_index < len(fits_list):
            data_next = load_fits_file(fits_list[next_index], 0, extra, reverse_flag=True)
        else:
            nchan = get_nchan_from_fits(fits_list[current_index])
            data_next = np.zeros((nchan, extra), dtype=np.float32)
        all_data.append(data_next)

    return np.hstack(all_data)


def extract_candidate_segment(fits_list, current_index, start_sample, DM):
    fits_file = fits_list[current_index]
    f_low, f_high, tsamp, freq_array, mjd_start = get_fits_freq_time_params(fits_file)
    DM_samples = calculate_dm_delay_samples(DM, f_low, f_high, tsamp)

    seg_start = start_sample - 5000
    if DM_samples < 20000:
        Dm_samples = 20000
    seg_end = start_sample + DM_samples+5000

    data = read_padded_samples(fits_list, current_index, seg_start, seg_end)
    ntime = data.shape[1]
    sample_indices = np.arange(seg_start, seg_end)
    time_array = (sample_indices - start_sample) * tsamp
    return data, freq_array, time_array, tsamp, mjd_start


# ================================
# Step 3: 数据预处理工具
# ================================
def normalise(data):
    data_ = data.copy()
    data_ /= np.mean(data_, axis=1)[:, None]
    data_ = (data_ - data_.min()) / (data_.max() - data_.min())
    return data_


def dedisperse(data, freq, tsamp, dm):
    nf, nt = data.shape
    delay_bins = (
        4.15 * dm * (freq.max() ** -2 - freq ** -2) * 1e3 / tsamp
    ).astype(np.int64)
    dedispersed = np.empty(data.shape, dtype=np.float32)
    for ii in range(nf):
        dedispersed[ii, :] = np.concatenate(
            [data[ii, -delay_bins[ii]:], data[ii, : -delay_bins[ii]]]
        )
    return dedispersed


def down_sample(data, f_down, t_down):
    shape = data.shape
    d_new_freq = np.empty((shape[0] // f_down, shape[1]))
    d_new = np.empty((shape[0] // f_down, shape[1] // t_down))
    for f in range(shape[0] // f_down):
        d_new_freq[f, :] = data[f * f_down:(f + 1) * f_down, :].sum(axis=0)
    for t in range(shape[1] // t_down):
        d_new[:, t] = d_new_freq[:, t * t_down:(t + 1) * t_down].sum(axis=1)
    return d_new


# ================================
# Step 4: 四图绘图函数
# ================================
def plot_frb_four_panel_summary(data_raw, data_dedispersed, freqs, times, tsamp,
                                 snr, dm, sample_id, fits_file, mjd_start,
                                 line_number=1, outname=None):
    from matplotlib import gridspec

    total_samples = data_raw.shape[1]
    center_index = 5000
    times = (np.arange(total_samples) - center_index) * tsamp

    t0 = mjd_start + (sample_id * tsamp) / 86400.0
    t_utc = Time(t0, format='mjd').to_datetime()
    t_bj = t_utc + timedelta(hours=8)

    # ========== 图像大小设置（全部正方形） ==========
    fig = plt.figure(figsize=(10, 10))
    gs = gridspec.GridSpec(2, 2)

    ax1 = fig.add_subplot(gs[0, 0])  # Raw waterfall
    ax2 = fig.add_subplot(gs[0, 1])  # Dedispersed waterfall
    ax3 = fig.add_subplot(gs[1, 0])  # Zoomed
    ax4 = fig.add_subplot(gs[1, 1])  # Pulse profile


    
    
    f_down_raw, t_down_raw = 16, 32
    rdw = down_sample(data_raw, f_down_raw, t_down_raw)
    freq_rds = np.linspace(freqs[0], freqs[-1], rdw.shape[0])
    time_rds = np.linspace(times[0], times[-1], rdw.shape[1])    
    dirdw = down_sample(data_dedispersed, f_down_raw, t_down_raw)

    mean_val = np.mean(rdw)
    std_val = np.std(rdw)
    vmin = mean_val - 2 * std_val
    vmax = mean_val + 2 * std_val
    # Raw Waterfall
    ax = ax1
    im = ax.pcolormesh(time_rds, freq_rds, rdw, cmap='viridis',
                       shading='auto', vmin=vmin, vmax=vmax)
    ax.set_title("Raw Waterfall (No Dedispersion)")
    ax.set_ylabel("Frequency (MHz)")
    ax.set_xlim(times[0], times[-1])

    # Dispersion lines
    f_ref = np.max(freqs)
    delay_ms = 4.148808e3 * dm * (1.0 / freqs**2 - 1.0 / f_ref**2)
    delay_s = delay_ms
    time_delay = delay_s.max() - delay_s.min()
    t0_line = 0.0
    t_line1 = t0_line - 0.125 * time_delay + delay_s
    t_line2 = t0_line + 0.125 * time_delay + delay_s
    ax.plot(t_line1, freqs, '--', color='red', linewidth=1)
    ax.plot(t_line2, freqs, '--', color='red', linewidth=1)
    ax.axvline(x=0, color='black', linestyle='--', linewidth=1)

    # 消色散后瀑布图与脉冲轮廓
    gs2 = gridspec.GridSpecFromSubplotSpec(2, 1, subplot_spec=gs[0, 1], height_ratios=[1, 5], hspace=0.0)
    ax2a = fig.add_subplot(gs2[0])  # 上：脉冲轮廓
    ax2b = fig.add_subplot(gs2[1], sharex=ax2a)  # 下：去色散瀑布

    # 上方：脉冲轮廓
    profile_dedisp = np.sum(dirdw, axis=0)
    ax2a.plot(time_rds, profile_dedisp, color='black')
    ax2a.set_title("Dedispersed Profile + Waterfall")
    ax2a.tick_params(labelbottom=False)  # 不显示x轴刻度
    ax2a.set_ylabel("Amplitude")

    # 下方：去色散瀑布图
    ax2b.pcolormesh(time_rds, freq_rds, dirdw, cmap='viridis',
                    shading='auto', vmin=vmin, vmax=vmax)
    ax2b.set_ylabel("Frequency (MHz)")
    ax2b.set_xlim(times[0], times[-1])
    ax2b.axvline(x=0.1, color='black', linestyle='--', linewidth=1)

    # Zoomed Waterfall
    cut_center = center_index
    time_window = 0.1 ####second 100ms width
    cut_width = int(time_window / tsamp) / 2
    
    start = cut_center - cut_width
    end = cut_center + cut_width
    sub_data = data_dedispersed[:, start:end]
    sub_times = times[start:end]
    
    
    t_rsl = tsamp  # time resolution
    fdown_zoom = 16
    if snr < 8:
        tdown_zoom = 8 * int(t_rsl / 4.9152e-05)
    elif 8 <= snr < 12:
        tdown_zoom = 4 * int(t_rsl / 4.9152e-05)
    elif 12 <= snr < 20:
        tdown_zoom = 2 * int(t_rsl / 4.9152e-05)
    else:
        tdown_zoom = 1 * int(t_rsl / 4.9152e-05)
    
    d_ds = down_sample(sub_data, fdown_zoom, tdown_zoom)
    freq_ds = freqs.reshape(-1, fdown_zoom).mean(axis=1)
    vminz = np.mean(d_ds) - 2 * np.std(d_ds)
    vmaxz = np.mean(d_ds) + 2 * np.std(d_ds)
    # Zoomed Waterfall
    ax = ax3
    ax.pcolormesh(
        np.linspace(sub_times[0], sub_times[-1], d_ds.shape[1]),
        freq_ds, d_ds,
        shading='auto', cmap='viridis', vmax=vmaxz, vmin=vminz
    )
    ax.set_title("Zoomed Waterfall (Downsampled)")
    ax.set_ylabel("Frequency (MHz)")
    ax.set_xlabel("Time (s)")

    # Pulse Profile
    profile = np.sum(sub_data, axis=0)
    ax = ax4
    ax.plot(sub_times, profile, color='black')
    ax.set_title("Pulse Profile")
    ax.set_xlabel("Time (s)")

    # 总标题框上移避免遮挡
    textstr = (
        f"Line: {line_number} | File: {os.path.basename(fits_file)}\n"
        f"Sample ID: {sample_id} | SNR: {snr:.2f} | DM: {dm:.2f}\n"
        f"MJD: {t0:.12f} | UTC: {t_utc.strftime('%Y-%m-%d %H:%M:%S')} |BJT: {t_bj.strftime('%Y-%m-%d %H:%M:%S')}\n"
        f"Delay: {time_delay:.4f} s"
    )
    fig.text(0.5, 0.99, textstr, fontsize=10, ha='center',
             va='top', bbox=dict(boxstyle="round", facecolor='lightcoral', edgecolor='black'))


    plt.savefig(outname)
    plt.close()
    print(f"Saved: {outname}")

